package com.basker.pisces.trace;

import org.springframework.core.NamedThreadLocal;
import org.springframework.lang.Nullable;

import com.basker.ttl.NamedTransmittableThreadLocal;

/**
 * {@link TraceContext}线程寄存器
 *
 * @author hangwen
 */
public class TraceContextHolder {
    private static final String NAME = TraceContextHolder.class.getSimpleName();

    private static final ThreadLocal<TraceContext> TL_HOLDER = new NamedThreadLocal<>(NAME);
    private static final ThreadLocal<TraceContext> TTL_HOLDER = new NamedTransmittableThreadLocal<>(NAME);

    public static TraceContext getTraceContext() {
        TraceContext tcx = TL_HOLDER.get();
        if (tcx == null) {
            tcx = TTL_HOLDER.get();
        }
        return tcx;
    }

    public static void remove() {
        TL_HOLDER.remove();
        TTL_HOLDER.remove();
    }

    public static void setTraceContext(@Nullable TraceContext tcx, boolean inheritable) {
        if (tcx == null) {
            return;
        }

        if (inheritable) {
            TTL_HOLDER.set(tcx);
            TL_HOLDER.remove();
        } else {
            TL_HOLDER.set(tcx);
            TTL_HOLDER.remove();
        }
    }
}
